import re
import csv
import pandas
from pyfaidx import Fasta
from .transcripts import read_transcripts
from .splicing import SplicingMaxEnt


def av_to_vcf(chrom: str, start: int, end: int, ref: int, alt: int, fasta: Fasta) -> [str, int, str, str]:
    chrom = 'chr' + re.sub(r'[Cc][Hh][Rr]', '', chrom)
    if chrom == 'chrMT':
        chrom = 'chrM'
    if ref == '-':
        pos = int(start)
        ref = fasta[chrom][pos-1:pos].seq  # fetch_seq(fasta=fasta, chrom=chrom, start=pos, end=pos)
        alt = f'{ref}{alt}'
    elif alt == '-':
        pos = int(start) - 1
        alt = fasta[chrom][pos-1:pos].seq  # fetch_seq(fasta=fasta, chrom=chrom, start=pos, end=pos)
        ref = f'{alt}{ref}'
    else:
        pos = int(start)
    return chrom, pos, ref, alt


def run_maxentpy(annovar_file: str, genome_file: str, refgene_file: str, gene_detail_columns: list[str], outfile: str):
    genome = Fasta(genome_file)
    transcripts = read_transcripts(refgene_file)
    reader = pandas.read_csv(annovar_file, sep="\t", dtype={'Chr': str})
    trans_dict = dict()
    rows = list()
    for row in reader.iloc:
        chrom, pos, ref, alt = av_to_vcf(row.Chr, row.Start, row.End, row.Ref, row.Alt, genome)
        transcript_names = set()
        for gene_detail_column in gene_detail_columns:
            gene_detail = row.get(gene_detail_column, '.')
            if isinstance(gene_detail, str) and gene_detail != ".":
                for detail in re.split(r',|;', gene_detail):
                    info = detail.split(':')
                    name = info[1] if len(info) == 5 else info[0]
                    transcript_names.add(name)
        result = None
        for transcript_name in transcript_names:
            trans_dict.setdefault(transcript_name, transcripts.get(transcript_name))
            transcript = trans_dict.get(transcript_name)
            if transcript:
                splicing = SplicingMaxEnt(chrom, pos, ref, alt, transcript, genome)
                if result:
                    if abs(splicing.maxentscore_var) > abs(result.maxentscore_var):
                        result = splicing
                    elif abs(splicing.maxentscore_var) == abs(result.maxentscore_var):
                        if splicing.maxentscore_alt > result.maxentscore_alt:
                            result = splicing
                else:
                    result = splicing
        if result:
            rows.append({
                'Chr': row.Chr, 'Start': row.Start, 'End': row.End, 'Ref': row.Ref, 'Alt': row.Alt,
                'Maxent_type': result.splice_type,
                'Maxent_pred': result.maxentpred,
                'Maxent_score_ref': result.maxentscore_ref,
                'Maxent_score_alt': result.maxentscore_alt,
                'Maxent_score_var': result.maxentscore_var,
                'Maxent_foldchange': result.maxentscore_foldchange
            })

    with open(outfile, 'w') as fo:
        writer = csv.DictWriter(fo, fieldnames=[
            'Chr',
            'Start',
            'End',
            'Ref',
            'Alt',
            'Maxent_type',
            'Maxent_pred',
            'Maxent_score_ref',
            'Maxent_score_alt',
            'Maxent_score_var',
            'Maxent_foldchange'
        ], delimiter='\t')
        writer.writeheader()
        writer.writerows(rows)
